"""
数据集脚本设置文件
包含模型配置、设备设置、训练参数等全局设置
用于统一管理各种模型和训练相关的配置参数
"""

import torch

# ==================== 模型配置 ====================
# Embedding模型名称 - 用于文本向量化
embedding_modelname = "Qwen/Qwen3-Embedding-0.6B"

# Reranker模型名称 - 用于重排序检索结果
reranker_modelname = "Qwen/Qwen3-Reranker-0.6B"

# 大语言模型名称 - 用于生成回答
llm_modelname = "Qwen/Qwen3-8B"

# 奖励模型名称 - 用于RLHF训练中评估回答质量
rewardmodel_modelname = "Qwen/Qwen3-0.6B"

# ==================== 设备配置 ====================
# 自动检测并设置计算设备（优先使用GPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ==================== 通用训练参数 ====================
# 随机种子 - 确保实验可重现
random_seed = 7

# 负样本文档数量 - 用于对比学习
num_negative_docs = 10

# ==================== Embedding模型参数 ====================
# 最大输入长度
embedding_max_length = 512
# 批处理大小
embedding_batch_size = 4
# 测试集比例
embedding_test_ratio = 0.2

# ==================== Reranker模型参数 ====================
# 最大输入长度
reranker_max_length = 512
# 批处理大小
reranker_batch_size = 4
# 测试集比例
reranker_test_ratio = 0.2

# ==================== LLM模型参数 ====================
# 最大输入长度
llm_max_length = 1024
# 批处理大小
llm_batch_size = 4
# 测试集比例
llm_test_ratio = 0.2

# ==================== 奖励模型参数 ====================
# 最大输入长度
rewardmodel_max_length = 1024
# 批处理大小
rewardmodel_batch_size = 4
# 测试集比例
rewardmodel_test_ratio = 0.2